"""
深度伪造检测数据集处理模块

主要功能：负责实时的数据加载、处理和增强
1. 数据集读取与划分（训练集/验证集/测试集）
2. 数据增强与预处理
3. 实现Mixup数据增强策略
"""

import os
import re
import glob
import copy
import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset
from albumentations import Compose, RandomBrightnessContrast,HorizontalFlip, FancyPCA, HueSaturationValue, OneOf, ToGray, ShiftScaleRotate, ImageCompression, PadIfNeeded, GaussNoise, GaussianBlur,Resize
from albumentations.pytorch.functional import img_to_tensor
import cv2
import logging

class ReadDataset():
    """
    数据集读取类
    功能：
    - 读取并解析数据集文件
    - 支持数据集重采样平衡
    - 可重复用于训练/验证/测试
    """
    def __init__(self,dataset_name):
        """
        初始化数据集读取器
        Args:
            dataset_name: 数据集名称（如'FF++'等）
        """
        oversample=True
        self.data = {
            "train": [],
            "val":[],
            "test": [],
        }
        self.labels = copy.deepcopy(self.data)
        self.path=f'datasets/{dataset_name}'
        
        # 读取数据集并记录样本统计信息
        self.read_txt(oversample=oversample)
        logging.info(f"fake data: {sum(self.labels['train'])}, real data: {len(self.labels['train'])-sum(self.labels['train'])}")

    def read_txt(self,oversample=False):
        """
        读取数据集txt文件
        Args:
            oversample: 是否对少数类别进行重采样
            
        实现细节：
        1. 分别读取训练/验证/测试集的真假样本
        2. 根据数据集类型设置平衡比例
        3. 支持真实样本重采样以平衡数据集
        """
        # 构建数据集文件路径
        dataset_files=[os.path.join(self.path,'test_fake.txt'),os.path.join(self.path,'test_real.txt'),
                        os.path.join(self.path,'val_fake.txt'),os.path.join(self.path,'val_real.txt'),
                        os.path.join(self.path,'train_fake.txt'),os.path.join(self.path,'train_real.txt')]
        
        # 根据数据集设置平衡比例
        if 'FF++' in self.path or 'DFDC-Preview' in self.path:
            balance_ratio=4
        elif 'Celeb-DF-v2' in self.path:
            balance_ratio = 6
        elif 'DFDC' in self.path:
            balance_ratio = 5
        else:
            balance_ratio = 1

        # 处理每个数据集文件
        for file in dataset_files:
            with open(file, "r") as f:
                lines = f.readlines()
            # 确定数据集类型
            if '/test_' in file:
                key='test'
            elif '/val_' in file:
                key='val'
            elif '/train_' in file:
                key = 'train'
                
            # 处理每一行数据
            for i in range(len(lines)):
                raw = re.sub("\s", "", lines[i]).split(",")
                paths = os.listdir(raw[1])  # 视频帧目录
                for row in paths:
                    path_dir = os.path.join(raw[1], row)
                    if os.path.isfile(path_dir) and '.png' in path_dir:
                        # 对训练集真实样本进行重采样
                        if oversample and 'train_real' in file:
                            for i in range(balance_ratio):
                                self.data[key].append(path_dir)
                                self.labels[key].append(int(raw[0]))
                        else:
                            self.data[key].append(path_dir)
                            self.labels[key].append(int(raw[0]))

    def get_dataset(self, mode):
        """获取指定模式的数据集"""
        return self.data[mode], self.labels[mode]

class MyDataset(Dataset):
    """
    自定义数据集类
    功能：
    - 实现数据加载和预处理
    - 提供丰富的数据增强方法
    """
    def __init__(self,
                 data,
                 label,
                 size=224,
                 normalize={"mean": [0.485, 0.456, 0.406],
                            "std": [0.229, 0.224, 0.225]},
                 test=False):
        """
        初始化数据集
        Args:
            data: 图像路径列表
            label: 标签列表
            size: 图像目标尺寸
            normalize: 归一化参数
            test: 是否为测试模式
        """
        super().__init__()
        self.size=size
        self.data=data
        self.label = label
        self.normalize = normalize
        self.aug=self.create_train_aug()
        self.transform=self.transform_all()
        self.test=test

    def create_train_aug(self):
        """
        创建训练数据增强流水线
        包含：
        - 图像压缩
        - 高斯噪声
        - 高斯模糊
        - 水平翻转
        - 亮度对比度调整
        - 灰度转换
        - 平移旋转缩放
        """
        return Compose([
            ImageCompression(quality_lower=60, quality_upper=100, p=0.5),  # 模拟压缩失真
            GaussNoise(p=0.1),  # 添加噪声
            GaussianBlur(blur_limit=3, p=0.05),  # 模糊处理
            HorizontalFlip(),  # 水平翻转
            PadIfNeeded(min_height=self.size, min_width=self.size, border_mode=cv2.BORDER_CONSTANT),
            OneOf([RandomBrightnessContrast(), FancyPCA(), HueSaturationValue()], p=0.7),  # 颜色增强
            ToGray(p=0.2),  # 灰度转换
            ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=10, border_mode=cv2.BORDER_CONSTANT, p=0.5),
        ])

    def transform_all(self):
        """创建基础变换：调整图像大小"""
        return Resize(p=1, height=self.size, width=self.size)

    def __getitem__(self, idx):
        """
        获取单个数据样本
        流程：
        1. 读取图像
        2. 应用基础变换
        3. 训练模式下应用数据增强
        4. 转换为张量并归一化
        """
        img = cv2.imread(self.data[idx], cv2.IMREAD_COLOR)
        data = self.transform(image=img)
        img = data["image"]
        if not self.test:
            data = self.aug(image=img)
            img = data["image"]

        img=img_to_tensor(img,self.normalize)
        return img,self.label[idx]

    def __len__(self):
        return len(self.data)

def mixup_data(x, y, alpha=0.5, use_cuda=False):
    """
    实现Mixup数据增强
    Args:
        x: 输入图像批次
        y: 标签批次
        alpha: Beta分布参数
        use_cuda: 是否使用GPU
    Returns:
        mixed_x: 混合后的图像
        y_a, y_b: 原始标签对
        lam: 混合权重
    """
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    batch_size = x.size()[0]
    if use_cuda:
        index = torch.randperm(batch_size).cuda()
    else:
        index = torch.randperm(batch_size)

    mixed_x = lam * x + (1 - lam) * x[index, :]
    y_a, y_b = y, y[index]
    return mixed_x, y_a, y_b, lam

def mixup_criterion(criterion, pred, y_a, y_b, lam):
    """
    计算Mixup损失
    Args:
        criterion: 损失函数
        pred: 模型预测
        y_a, y_b: 标签对
        lam: 混合权重
    Returns:
        混合后的损失值
    """
    a=criterion(pred, y_a)
    b=criterion(pred, y_b)
    losses=[]
    try:
        for i in range(len(a)):
            losses.append(lam * a[i]  + (1 - lam) * b[i])
    except:
        return lam * a  + (1 - lam) * b
    return losses
